# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
from functools import partial

import paddle
from paddlenlp.datasets import load_dataset, MapDataset
from paddlenlp.transformers import AutoTokenizer
from paddlenlp.metrics import SpanEvaluator
from paddlenlp.utils.log import logger

from model import UIE
from utils import convert_example, reader, unify_prompt_name, get_relation_type_dict, create_data_loader


@paddle.no_grad()
def evaluate(model, metric, data_loader):
    """
    Given a dataset, it evals model and computes the metric.
    Args:
        model(obj:`paddle.nn.Layer`): A model to classify texts.
        metric(obj:`paddle.metric.Metric`): The evaluation metric.
        data_loader(obj:`paddle.io.DataLoader`): The dataset loader which generates batches.
    """
    model.eval()
    metric.reset()
    for batch in data_loader:
        input_ids, token_type_ids, att_mask, pos_ids, start_ids, end_ids = batch
        start_prob, end_prob = model(input_ids, token_type_ids, att_mask,
                                     pos_ids)
        start_ids = paddle.cast(start_ids, 'float32')
        end_ids = paddle.cast(end_ids, 'float32')
        # todo 举例说明
        # metric = SpanEvaluator()
        # 在metric.compute(start_prob, end_prob, start_ids, end_ids)方法中，
        #  todo step1：对每条样本，筛选出其预测得分集里大于阈值的index。
        #   假设某条样本(maxlength=30)输出 “状态 = 是某个实体的start” 的预测得分集 = [
        #             0.00000031, 0.00000129, 0.00000024, 0.00000031, 0.50000042, 0.00070048,
        #             0.00000877, 0.00002038, 0.00000138, 0.00001625, 0.00000029, 0.00008086,
        #             0.00000202, 0.00000990, 0.00000406, 0.00016828, 0.00001869, 0.00021678,
        #             0.00000937, 0.00046841, 0.00002430, 0.00000023, 0.00006322, 0.00000595,
        #             0.00000095, 0.00000166, 0.00000030, 0.00000033, 0.00000023, 0.00000073]
        #   则得分大于0.5的index=[4]，即预测当前样本 index=4 处可能是一个实体的开始位置。
        #   依次对batch中每个样本的预测得分集做以上处理，得到每个样本的pre-start和pre-end。综合按batch输出，即：
        # pre-start = [[4],  [],  [14],  [4],  [8],   [],  [5, 8], [5, 8]]  # 可以看到当前批次的第7，8条样本中各有2个位置处的得分 > 阈值0.5。
        # pre_end =   [[7],  [],   [],   [7],  [9],   [],   [6],   [6, 9]]
        #
        # gold_start= [[4], [14], [14],  [4],  [8],  [13],   [5],    [8]]
        # glod_end =  [[7], [14], [14],  [7],  [8],  [14],   [6],    [9]]

        #  todo step2：沿文本序列方向组合(pre_start,pre_end)、(gold_start,gold_end)，得到
        # pre： (4,7),   (),     (),  (4,7),(8,9),   (),  (5,6), ((5,6),(8,9))
        # gold：(4,7),(14,14),(14,14),(4,7),(8,8),(13,14),(5,6),     (8,9)
        #
        #  todo step3：pre、gold对应组合求与&运算，得到：
        # result: 1,     0,      0,     1,    1,     0,     1,         1
        # pre： (4,7),   (),     (),  (4,7),(8,9),   (),  (5,6), ((5,6),(8,9))
        # gold：(4,7),(14,14),(14,14),(4,7),(8,8),(13,14),(5,6),     (8,9)

        # todo step4：统计预测组pre、真实组gold、正确组correct的数量
        # correct: 1,    0,      0,     1,    1,     0,     1,         1           --num_correct=5
        # pre： (4,7),   (),     (),  (4,7),(8,9),   (),  (5,6), ((5,6),(8,9))     --num_infer=6
        # gold：(4,7),(14,14),(14,14),(4,7),(8,8),(13,14),(5,6),     (8,9)         --num_label=8

        # todo step5：计算precious、recall、F1
        # precious = num_correct/num_infer
        # recall=num_correct/num_label
        # F1=2*precious*recall/(precious+recall)
        num_correct, num_infer, num_label = metric.compute(
            start_prob, end_prob, start_ids, end_ids)
        metric.update(num_correct, num_infer, num_label)
    precision, recall, f1 = metric.accumulate()
    model.train()
    return precision, recall, f1


def do_eval():
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = UIE.from_pretrained(args.model_path)

    test_ds = load_dataset(reader,
                           data_path=args.test_path,
                           max_seq_len=args.max_seq_len,
                           lazy=False)
    class_dict = {}
    relation_data = []
    if args.debug:
        for data in test_ds:
            class_name = unify_prompt_name(data['prompt'])
            # Only positive examples are evaluated in debug mode
            if len(data['result_list']) != 0:
                if "的" not in data['prompt']:
                    class_dict.setdefault(class_name, []).append(data)
                else:
                    relation_data.append((data['prompt'], data))
        relation_type_dict = get_relation_type_dict(relation_data)
    else:
        class_dict["all_classes"] = test_ds

    trans_fn = partial(convert_example,
                       tokenizer=tokenizer,
                       max_seq_len=args.max_seq_len)

    for key in class_dict.keys():
        if args.debug:
            test_ds = MapDataset(class_dict[key])
        else:
            test_ds = class_dict[key]

        test_data_loader = create_data_loader(test_ds,
                                              mode="test.txt",
                                              batch_size=args.batch_size,
                                              trans_fn=trans_fn)

        metric = SpanEvaluator()
        precision, recall, f1 = evaluate(model, metric, test_data_loader)
        logger.info("-----------------------------")
        logger.info("Class Name: %s" % key)
        logger.info("Evaluation Precision: %.5f | Recall: %.5f | F1: %.5f" %
                    (precision, recall, f1))

    if args.debug and len(relation_type_dict.keys()) != 0:
        for key in relation_type_dict.keys():
            test_ds = MapDataset(relation_type_dict[key])

            test_data_loader = create_data_loader(test_ds,
                                                  mode="test.txt",
                                                  batch_size=args.batch_size,
                                                  trans_fn=trans_fn)

            metric = SpanEvaluator()
            precision, recall, f1 = evaluate(model, metric, test_data_loader)
            logger.info("-----------------------------")
            logger.info("Class Name: X的%s" % key)
            logger.info("Evaluation Precision: %.5f | Recall: %.5f | F1: %.5f" %
                        (precision, recall, f1))


if __name__ == "__main__":
    # yapf: disable
    parser = argparse.ArgumentParser()

    parser.add_argument("--model_path", type=str, default=None, help="The path of saved model that you want to load.")
    parser.add_argument("--test_path", type=str, default=None, help="The path of test.txt set.")
    parser.add_argument("--batch_size", type=int, default=16, help="Batch size per GPU/CPU for training.")
    # todo 默认是512，debug时可以设置小点，方便看矩阵里的值，比如30。
    # parser.add_argument("--max_seq_len", type=int, default=512, help="The maximum total input sequence length after tokenization.")
    parser.add_argument("--max_seq_len", type=int, default=30, help="The maximum total input sequence length after tokenization.")
    parser.add_argument("--debug", action='store_true', help="Precision, recall and F1 score are calculated for each class separately if this option is enabled.")

    args = parser.parse_args()
    # yapf: enable

    do_eval()